import numpy as np
from PIL import Image
import matplotlib.pyplot as plt


class SVD:
    def __init__(self, img_path):
        with Image.open(img_path) as img:
            img = np.asarray(img.convert('L'))
        self.U, self.Sigma, self.VT = np.linalg.svd(img)

    def compress_img(self, k: "# singular value") -> "img":
        return self.U[:, :k] @ np.diag(self.Sigma[:k]) @ self.VT[:k, :]


if __name__ == '__main__':
    model = SVD('./Lenna.png')
    result = [
        Image.fromarray(model.compress_img(i))
        for i in [1, 10, 20, 50, 100, 500]
    ]
    img_array0 = np.array(result[0])
    img_array1 = np.array(result[1])
    img_array2 = np.array(result[2])
    img_array3 = np.array(result[3])
    img_array4 = np.array(result[4])
    img_array5 = np.array(result[5])
    fig, axs = plt.subplots(nrows=2, ncols=3, figsize=(13, 8))
    axs[0][0].imshow(img_array0)
    axs[0][1].imshow(img_array1)
    axs[0][2].imshow(img_array2)
    axs[1][0].imshow(img_array3)
    axs[1][1].imshow(img_array4)
    axs[1][2].imshow(img_array5)

    plt.show()
